Skip to content

Optimize JitFunction direct CallState dispatch - #668

Closed
yhl-amd wants to merge 2 commits into
mainfrom
yhl/direct-callstate-hotpath
Closed

Optimize JitFunction direct CallState dispatch#668
yhl-amd wants to merge 2 commits into
mainfrom
yhl/direct-callstate-hotpath

Conversation

@yhl-amd

@yhl-amd yhl-amd commented Jun 9, 2026

Copy link
Copy Markdown

Summary

  • avoid wrapping runtime arguments only to build the JIT cache key
  • cache a direct CallState hot path for JitFunctions without self, Constexpr, or type parameters
  • support kwargs/default reconstruction when using the direct call path

Context

This is the FlyDSL side of the DeepSeek-V4-Pro non-DPA 1k/1k c=256 regression fix. The regression was dominated by host-side launcher overhead in prefill attention-adjacent FlyDSL launchers, not by generated GPU kernel quality.

Paired with the AITER qk launcher patch, local ATOM E2E recovered most of the regression:

case output tok/s TPOT
old good 4250.48 58.04 ms
bad new 3907.11 63.32 ms
paired fix 4227.06 58.30 ms

Validation

  • python -m py_compile python/flydsl/compiler/jit_function.py tests/unit/test_jit_direct_call_state.py
  • PYTHONPATH=python python -m pytest tests/unit/test_jit_direct_call_state.py -q is blocked in my local source tree because it is not built and flydsl._mlir is absent; CI should run it after the normal build step.
  • ATOM DeepSeek-V4-Pro non-DPA c=256 E2E with paired AITER patch: 4227.06 output tok/s, 58.30 ms TPOT.

@yhl-amd yhl-amd mentioned this pull request Jun 9, 2026
Merged
@yhl-amd
yhl-amd marked this pull request as ready for review June 9, 2026 01:47
@yhl-amd
yhl-amd force-pushed the yhl/direct-callstate-hotpath branch from f82134a to 610d9f9 Compare June 9, 2026 09:17
@yhl-amd

yhl-amd commented Jun 9, 2026

Copy link
Copy Markdown
Author

Added a FlyDSL-only hotpath replay benchmark in commit 610d9f94:

benchmarks/dsv4_hotpath_repro/
  bench_flydsl_hotpath.py
  run_matrix.sh
  README.md
  baselines.md

Purpose: reproduce the host-side FlyDSL launch overhead that showed up in the DeepSeek-V4-Pro non-DPA 1k/1k c=256 regression, without requiring AITER, ATOM, model weights, or a serving stack. The benchmark uses synthetic FlyDSL kernels with AITER-like launcher signatures and rank-local trace call counts:

qk_norm_rope_quant_like:      81 calls/window
fused_compress_attn_like:     80 calls/window
hca_compress_forward_like:    41 calls/window
hca_norm_rope_scatter_like:   41 calls/window

How to run a quick comparison from this PR checkout:

benchmarks/dsv4_hotpath_repro/run_matrix.sh \
  --original /path/to/original-flydsl \
  --fixed /path/to/fixed-flydsl \
  -- --windows 1 --warmup-windows 0 --gpu-event-calls 2

--fixed defaults to the current checkout, so after building this branch you can also run:

benchmarks/dsv4_hotpath_repro/run_matrix.sh \
  --original /path/to/original-flydsl \
  -- --windows 1 --warmup-windows 0 --gpu-event-calls 2

The helper detects these layouts for each path:

/path/to/FlyDSL/build-fly/python_packages
/path/to/FlyDSL/build/python_packages
/path/to/FlyDSL/python
/path/to/FlyDSL

For a longer replay closer to the original c=256 profile, drop the smoke args and use the defaults (--windows 16).

Local smoke result on gfx950 in the yhl_dev ROCm container:

mixed jit host time:
  original FlyDSL 0.2.0-pristine:      104.23 us/call
  fixed directraw-kw snapshot:           91.36 us/call
  speedup:                                1.14x

mixed compiled-position call time:
  original:                              10.09 us/call
  fixed:                                 10.01 us/call

qk_norm_rope_quant_like jit host time:
  original:                              68.68 us/call
  fixed:                                 64.89 us/call
  speedup:                                1.06x

attention-adjacent jit host time:
  fused_compress:                       152.45 -> 129.27 us/call  (1.18x)
  hca_compress:                         106.40 ->  90.80 us/call  (1.17x)
  hca_scatter:                           91.67 ->  78.31 us/call  (1.17x)

Interpretation: compiled_positional stays around ~10 us/call in both stacks, while the warm @flyc.jit path is much higher and improves with the hotpath fix. That isolates the regression to host-side JIT/cache-key/TensorAdaptor/CallState overhead, not GPU kernel math. This is intentionally not MoE-only; qk norm quant and the attention-adjacent launchers are included because those were the high-frequency calls in the original c=256 trace.

@sjfeng1999 sjfeng1999 closed this Jun 11, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@yhl-amd @sjfeng1999